package protocol

import (
	"bytes"
	"encoding/json"
	"fmt"
)

// JSONRPCMessage JSON-RPC 2.0消息
type JSONRPCMessage struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      interface{}     `json:"id,omitempty"`
	Method  string          `json:"method,omitempty"`
	Params  json.RawMessage `json:"params,omitempty"`
	Result  json.RawMessage `json:"result,omitempty"`
	Error   *RPCError       `json:"error,omitempty"`
}

// RPCError JSON-RPC错误
type RPCError struct {
	Code    int         `json:"code"`
	Message string      `json:"message"`
	Data    interface{} `json:"data,omitempty"`
}

// 错误代码常量
const (
	ParseError     = -32700
	InvalidRequest = -32600
	MethodNotFound = -32601
	InvalidParams  = -32602
	InternalError  = -32603

	// MCP特定错误
	ToolNotFound  = -32001
	ToolTimeout   = -32002
	UserCancelled = -32003
	GUIError      = -32004
)

// NewRequest 创建请求消息
func NewRequest(id interface{}, method string, params interface{}) (*JSONRPCMessage, error) {
	var paramsJSON json.RawMessage
	if params != nil {
		buffer := &bytes.Buffer{}
		encoder := json.NewEncoder(buffer)
		encoder.SetEscapeHTML(false)
		if err := encoder.Encode(params); err != nil {
			return nil, fmt.Errorf("failed to marshal params: %w", err)
		}
		paramsJSON = bytes.TrimRight(buffer.Bytes(), "\n")
	}

	return &JSONRPCMessage{
		JSONRPC: "2.0",
		ID:      id,
		Method:  method,
		Params:  paramsJSON,
	}, nil
}

// NewResponse 创建响应消息
func NewResponse(id interface{}, result interface{}) (*JSONRPCMessage, error) {
	var resultJSON json.RawMessage
	if result != nil {
		// 使用自定义编码器，不转义Unicode字符
		buffer := &bytes.Buffer{}
		encoder := json.NewEncoder(buffer)
		encoder.SetEscapeHTML(false)
		if err := encoder.Encode(result); err != nil {
			return nil, fmt.Errorf("failed to marshal result: %w", err)
		}
		// 移除末尾的换行符
		resultJSON = bytes.TrimRight(buffer.Bytes(), "\n")
	}

	return &JSONRPCMessage{
		JSONRPC: "2.0",
		ID:      id,
		Result:  resultJSON,
	}, nil
}

// NewErrorResponse 创建错误响应
func NewErrorResponse(id interface{}, code int, message string, data interface{}) *JSONRPCMessage {
	return &JSONRPCMessage{
		JSONRPC: "2.0",
		ID:      id,
		Error: &RPCError{
			Code:    code,
			Message: message,
			Data:    data,
		},
	}
}

// NewNotification 创建通知消息（无ID）
func NewNotification(method string, params interface{}) (*JSONRPCMessage, error) {
	var paramsJSON json.RawMessage
	if params != nil {
		buffer := &bytes.Buffer{}
		encoder := json.NewEncoder(buffer)
		encoder.SetEscapeHTML(false)
		if err := encoder.Encode(params); err != nil {
			return nil, fmt.Errorf("failed to marshal params: %w", err)
		}
		paramsJSON = bytes.TrimRight(buffer.Bytes(), "\n")
	}

	return &JSONRPCMessage{
		JSONRPC: "2.0",
		Method:  method,
		Params:  paramsJSON,
	}, nil
}

// Marshal 序列化消息
func (m *JSONRPCMessage) Marshal() ([]byte, error) {
	buffer := &bytes.Buffer{}
	encoder := json.NewEncoder(buffer)
	encoder.SetEscapeHTML(false)
	if err := encoder.Encode(m); err != nil {
		return nil, err
	}
	return bytes.TrimRight(buffer.Bytes(), "\n"), nil
}

// Unmarshal 反序列化消息
func Unmarshal(data []byte) (*JSONRPCMessage, error) {
	var msg JSONRPCMessage
	if err := json.Unmarshal(data, &msg); err != nil {
		return nil, err
	}

	// 验证JSONRPC版本
	if msg.JSONRPC != "2.0" {
		return nil, fmt.Errorf("invalid jsonrpc version: %s", msg.JSONRPC)
	}

	return &msg, nil
}

// IsRequest 判断是否为请求
func (m *JSONRPCMessage) IsRequest() bool {
	return m.Method != "" && m.ID != nil
}

// IsNotification 判断是否为通知
func (m *JSONRPCMessage) IsNotification() bool {
	return m.Method != "" && m.ID == nil
}

// IsResponse 判断是否为响应
func (m *JSONRPCMessage) IsResponse() bool {
	return m.Method == "" && (m.Result != nil || m.Error != nil)
}
